/*
https://leetcode-cn.com/problems/sum-of-distances-in-tree/solution/shu-zhong-ju-chi-zhi-he-by-leetcode-solution/
 */
import java.util.Arrays;

public class Solution834 {
    public int[] sumOfDistancesInTree(int n, int[][] edges) {
        int[] first=new int[n];
        Arrays.fill(first,-1);
        int[] next=new int[2*n];
        int[] go=new int[2*n];
        int c=0;
        for (int[] edge : edges) {
            next[c]=first[edge[0]];
            first[edge[0]]=c;
            go[c]=edge[1];
            c++;
            next[c]=first[edge[1]];
            first[edge[1]]=c;
            go[c]=edge[0];
            c++;
        }
        int[] dp=new int[n];
        int[] sz=new int[n];
        int[] ans=new int[n];
        dfs(0,-1,dp,sz,first,next,go);
        dfs2(0,-1,dp,sz,ans,first,next,go);
        return ans;
    }

    private void dfs2(int k, int f, int[] dp, int[] sz, int[] ans, int[] first, int[] next, int[] go) {
        ans[k]=dp[k];
        int p=first[k];
        while (p!=-1){
            if (go[p]!=f){
                int a=dp[k],b=sz[k],c=dp[go[p]],d=sz[go[p]];
                dp[k]-=dp[go[p]]+sz[go[p]];
                sz[k]-=sz[go[p]];
                dp[go[p]]+=dp[k]+sz[k];
                sz[go[p]]+=sz[k];
                dfs2(go[p],k,dp,sz,ans,first,next,go);
                dp[k]=a;
                sz[k]=b;
                dp[go[p]]=c;
                sz[go[p]]=d;
            }
            p=next[p];
        }
    }

    private void dfs(int k, int f, int[] dp, int[] sz, int[] first, int[] next, int[] go) {
        int p=first[k];
        sz[k]=1;
        while (p!=-1){
            if (go[p]!=f){
                dfs(go[p],k,dp,sz,first,next,go);
                sz[k]+=sz[go[p]];
                dp[k]+=dp[go[p]]+sz[go[p]];
            }
            p=next[p];
        }
    }

    public static void main(String[] args) {
        System.out.println(Arrays.toString(new Solution834().sumOfDistancesInTree(6, new int[][]{{0, 1}, {0, 2}, {2, 3}, {2, 4}, {2, 5}})));
    }
}
